#include "ObjMeshLoader.hpp"
#include <assimp/Importer.hpp>
#include <assimp/postprocess.h>
#include <assimp/scene.h>
#include <iostream>

namespace Model {

void ObjMeshLoader::ProcessMesh(aiMesh *mesh, const aiScene *scene) {

  std::vector<float> buffer;
  std::vector<VertexAttribFormat> format;
  std::vector<uint32_t> indices;

  for (uint32_t i = 0; i < mesh->mNumVertices; i++) {
    uint8_t sum_comps = 0;
    if (mesh->mTextureCoords[0]) {
      buffer.push_back(mesh->mTextureCoords[0][i].x);
      buffer.push_back(mesh->mTextureCoords[0][i].y);
      if (i < 1) {
        format.push_back({DataType::FLOAT32, AttribType::TexCoords, sum_comps});
        sum_comps += 2;
      }
    }

    if (i < 1) {
      format.push_back({DataType::FLOAT32, AttribType::Pos, sum_comps});
      sum_comps += 3;
    }
    buffer.push_back(mesh->mVertices[i].x);
    buffer.push_back(mesh->mVertices[i].y);
    buffer.push_back(mesh->mVertices[i].z);

    if (i < 1) {
      format.push_back({DataType::FLOAT32, AttribType::Norm, sum_comps});
      sum_comps += 3;
    }
    buffer.push_back(mesh->mNormals[i].x);
    buffer.push_back(mesh->mNormals[i].y);
    buffer.push_back(mesh->mNormals[i].z);

    if (mesh->HasTangentsAndBitangents()) {
      if (i < 1) {
        format.push_back({DataType::FLOAT32, AttribType::BiTan, sum_comps});
        sum_comps += 3;
        format.push_back({DataType::FLOAT32, AttribType::Tangents, sum_comps});
        sum_comps += 3;
      }
      buffer.push_back(mesh->mBitangents[i].x);
      buffer.push_back(mesh->mBitangents[i].y);
      buffer.push_back(mesh->mBitangents[i].z);
      buffer.push_back(mesh->mTangents[i].x);
      buffer.push_back(mesh->mTangents[i].y);
      buffer.push_back(mesh->mTangents[i].z);
    }

    if (mesh->HasVertexColors(0)) {
      if (i < 1) {
        format.push_back({DataType::FLOAT32, AttribType::VCols, sum_comps});
        sum_comps += 4;
      }
      buffer.push_back(mesh->mColors[0][i].r);
      buffer.push_back(mesh->mColors[0][i].g);
      buffer.push_back(mesh->mColors[0][i].b);
      buffer.push_back(mesh->mColors[0][i].a);
    }
  }

  for (uint32_t i = 0; i < mesh->mNumFaces; i++) {
    aiFace face = mesh->mFaces[i];
    for (uint32_t j = 0; j < face.mNumIndices; j++) {
      indices.push_back(face.mIndices[j]);
    }
  }

  auto meshData = MeshData::CreateMeshData();
  meshData->Assign(buffer, format);
  meshData->AssignIndex(indices);
  GetMeshes().push_back(meshData);
}

void ObjMeshLoader::ProcessNode(aiNode *node, const aiScene *scene) {
  for (uint32_t i = 0; i < node->mNumMeshes; i++) {
    aiMesh *mesh = scene->mMeshes[node->mMeshes[i]];
    ProcessMesh(mesh, scene);
  }

  // After we've processed all of the meshes (if any) we then recursively process each of
  // the children nodes
  for (uint32_t i = 0; i < node->mNumChildren; i++) {
    ProcessNode(node->mChildren[i], scene);
  }
}

ObjMeshLoader::ObjMeshLoader(const std::string &fileName) : _fileName(fileName) {
  Load();
}

ObjMeshLoader::~ObjMeshLoader() {
  _fileName.clear();
  for (auto &val : _meshes) {
    val.reset();
  }
  _meshes.clear();
}

ObjMeshLoader::MeshRefs &ObjMeshLoader::GetMeshes() { return _meshes; }

void ObjMeshLoader::Load() {
  Assimp::Importer importer;
  const aiScene *scene = importer.ReadFile(
      _fileName, aiProcess_Triangulate | aiProcess_FlipUVs | aiProcess_CalcTangentSpace);

  if (!scene || scene->mFlags == AI_SCENE_FLAGS_INCOMPLETE || !scene->mRootNode) {
    std::cout << "ERROR::ASSIMP:: " << importer.GetErrorString() << std::endl;
    return;
  }

  ProcessNode(scene->mRootNode, scene);
}

} // namespace Model